#!/usr/bin/env python
import dpdata
import matplotlib.pyplot as plt
import numpy as np

# 读取训练数据点
training_systems = dpdata.LabeledSystem("/home/zxg/be-cu_dp_course/dpgen/run/deepmd_file/Be16_hcp0.96_1.04/deepmd", fmt = "deepmd/npy")
predict = training_systems.predict("/home/zxg/be-cu_dp_course/dpgen/run/iter.000000/00.train/000/frozen_model.pb")

# 提取训练和预测的力分量
train_forces = training_systems["forces"]
predict_forces = predict["forces"]
train_energies = training_systems["energies"]
predict_energies = predict["energies"]

# 提取 x、y、z 轴力的分量
train_forces_x = [force[0] for sublist in train_forces for force in sublist]
train_forces_y = [force[1] for sublist in train_forces for force in sublist]
train_forces_z = [force[2] for sublist in train_forces for force in sublist]

predict_forces_x = [force[0] for sublist in predict_forces for force in sublist]
predict_forces_y = [force[1] for sublist in predict_forces for force in sublist]
predict_forces_z = [force[2] for sublist in predict_forces for force in sublist]

# 计算绝对差值
energydiff = np.abs(np.array(train_energies) - np.array(predict_energies))
forcediff_x = np.abs(np.array(train_forces_x) - np.array(predict_forces_x))
forcediff_y = np.abs(np.array(train_forces_y) - np.array(predict_forces_y))
forcediff_z = np.abs(np.array(train_forces_z) - np.array(predict_forces_z))

# 使用 np.histogram 绘制直方图
hist_energy, bin_edges= np.histogram(energydiff, bins=100)
hist_x, bin_edges  = np.histogram(forcediff_x, bins=100)
hist_y, bin_edges  = np.histogram(forcediff_y, bins=100)
hist_z, bin_edges  = np.histogram(forcediff_z, bins=100)

# 创建一个figure
plt.figure(figsize=(12, 9))  # 设置figure的大小

# 创建子图来绘制直方图
plt.subplot(2, 2, 1)
plt.bar(bin_edges[:-1], hist_energy, width=bin_edges[1] - bin_edges[0], color='blue', alpha=0.7)
plt.xlabel('Absolute Difference Energy (|DFT - DeepMD|)')
plt.ylabel('Number of Data Points')

# 创建子图来绘制 x 轴力的直方图
plt.subplot(2, 2, 2)
plt.bar(bin_edges[:-1], hist_x, width=bin_edges[1] - bin_edges[0], color='green', alpha=0.7)
plt.xlabel('Absolute Difference forces x (|DFT - DeepMD|)')
plt.ylabel('Number of Data Points')

# 创建子图来绘制 y 轴力的直方图
plt.subplot(2, 2, 3)
plt.bar(bin_edges[:-1], hist_y, width=bin_edges[1] - bin_edges[0], color='red', alpha=0.7)
plt.xlabel('Absolute Difference forces y (|DFT - DeepMD|)')
plt.ylabel('Number of Data Points')

# 创建子图来绘制 z 轴力的直方图
plt.subplot(2, 2, 4)
plt.bar(bin_edges[:-1], hist_z, width=bin_edges[1] - bin_edges[0], color='blue', alpha=0.7)
plt.xlabel('Absolute Difference forces z (|DFT - DeepMD|)')
plt.ylabel('Number of Data Points')

# 设置图片的分辨率和保存路径
plt.savefig('inter00_f_e_diff_histogram.png')